from PIL import Image
import os

def cut_and_save_subimages_for_folder(folder_path, subimage_count, save_folder):
    # 遍历指定文件夹中的所有图像文件
    for filename in os.listdir(folder_path):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            img_path = os.path.join(folder_path, filename)
            cut_and_save_subimages(img_path, subimage_count, save_folder)

def cut_and_save_subimages(img_path, subimage_count, save_folder, traj):
    img = Image.open(img_path)
    width, height = img.size
    sub_width, sub_height = 300, 300

    # 确保保存文件夹存在
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)

    base_name = os.path.splitext(os.path.basename(img_path))[0]

    # 定义横向和纵向的偏移量
    offset_horizontal = 40  # 横向偏移量
    offset_vertical = 40    # 纵向偏移量

    if traj == 'horizontal' or traj == 'all':
        # 横向轨迹切割
        for i in range(subimage_count):
            left = min(i * (width - sub_width) // (subimage_count - 1), width - sub_width)
            upper = max((height // 2) - (sub_height // 2) - offset_vertical, 0)
            box = (left, upper, left + sub_width, upper + sub_height)
            sub_img = img.crop(box)
            sub_img.save(os.path.join(save_folder, f"{base_name}_horizontal_{subimage_count}_{i + 1}.jpg"))

    if traj == 'vertical' or traj == 'all':
        # 纵向轨迹切割
        for i in range(subimage_count):
            upper = min(i * (height - sub_height) // (subimage_count - 1), height - sub_height)
            left = max((width // 2) - (sub_width // 2) - offset_horizontal, 0)
            box = (left, upper, left + sub_width, upper + sub_height)
            sub_img = img.crop(box)
            sub_img.save(os.path.join(save_folder, f"{base_name}_vertical_{subimage_count}_{i + 1}.jpg"))

    if traj == 'diagonal' or traj == 'all':
        # 斜向轨迹切割
        for i in range(subimage_count):
            left = i * (width - sub_width) // (subimage_count - 1)
            upper = i * (height - sub_height) // (subimage_count - 1)
            left = max(min(left - offset_horizontal, width - sub_width), 0)
            upper = max(min(upper - offset_vertical, height - sub_height), 0)
            if left >= width - sub_width or upper >= height - sub_height:
                break
            box = (left, upper, left + sub_width, upper + sub_height)
            sub_img = img.crop(box)
            sub_img.save(os.path.join(save_folder, f"{base_name}_diagonal_{subimage_count}_{i + 1}.jpg"))


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description="执行图像裁切操作")
    parser.add_argument("--scene", type=str)
    parser.add_argument("--traj", type=str)
    args = parser.parse_args()

    img_path = f'../public/images/scenes/{args.scene}'
    subimage_count = 100

    save_folder = r'./output_images'
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)

    # 删除原来的文件
    delete_files = os.listdir(save_folder)
    for item in delete_files:
        file_path = os.path.join(save_folder, item)
        os.remove(file_path)
    
    cut_and_save_subimages(img_path, subimage_count, save_folder, args.traj)
